// SPDX-License-Identifier: GPL-2.0
#include "vmlinux.h"
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>

/*
 * 最简单：使用 fentry 钩住内核中的 bpf_handle_try_charge_memcg。
 * 内核里你的定义：
 *   int bpf_handle_try_charge_memcg(void);
 *
 * fentry 程序原型：int func(struct pt_regs *ctx)
 * 用 BPF_PROG 宏也可，但这里直接写。
 */

char LICENSE[] SEC("license") = "GPL";

#if 0
// 如果需要内核 5.15+ 更好的 CO-RE 支持，保持包含 vmlinux.h。
SEC("fentry/bpf_handle_try_charge_memcg")
int BPF_PROG(on_bpf_handle_try_charge_memcg)
{
    // bpf_printk 输出到 trace_pipe
    bpf_printk("eBPF: bpf_handle_try_charge_memcg was called\\n");
    return 0;
}

/*
 * 也可以加一个 fexit 示例（可选）。
 */
SEC("fexit/bpf_handle_try_charge_memcg")
int BPF_PROG(on_bpf_handle_try_charge_memcg_exit, struct handle_try_charge_memcg_s *hcm, int ret)
{
    hcm->nr_pages = 100;
    bpf_printk("eBPF: bpf_handle_try_charge_memcg exit, ret=%d\\n", ret);
    return 0;
}

extern void bpf_handle_try_charge_memcg_s_set(struct handle_try_charge_memcg_s *dst, struct handle_try_charge_memcg_s *src) __ksym;

// 对应：SEC("fmod_ret/bpf_handle_try_charge_memcg")
SEC("fmod_ret/bpf_handle_try_charge_memcg")
u64 BPF_PROG(mod_ret_bpf_handle_try_charge_memcg, struct handle_try_charge_memcg_s *hcm)
{
    struct handle_try_charge_memcg_s new;

    //new.nr_pages = 100;
    //bpf_handle_try_charge_memcg_s_set(hcm, &new);

    bpf_printk("eBPF: mod_ret_bpf_handle_try_charge_memcg");

    return 1;
}
#endif

/* eBPF 程序实现 */
SEC("struct_ops/handle_try_charge")
int handle_try_charge_impl(struct handle_try_charge_memcg_s *hcm)
{
    /* 修改 nr_pages 为 100 */
    hcm->nr_pages = 100;
    
    /* 可选：添加一些调试日志 */
    bpf_printk("eBPF: Modified nr_pages to %lu, gfp_mask: 0x%x\n", 
               hcm->nr_pages, hcm->gfp_mask);
    
    /* 返回 0 表示继续正常执行 */
    return 1;
}

/* 定义 struct ops map */
SEC(".struct_ops")
struct try_charge_memcg_ops try_charge_ops = {
    .handle_try_charge = (void *)handle_try_charge_impl,
};

char _license[] SEC("license") = "GPL";
